import numpy as np
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
import matplotlib.pyplot as plt
import torch




# 定义函数计算混淆矩阵
def calculate_confusion_matrix(dataloader, model, device, num_classes):
    model.eval()  # 设置模型为评估模式
    all_labels = []
    all_predictions = []

    with torch.no_grad():
        for (rgb_images, depth_images), labels in dataloader:
            rgb_images = rgb_images.to(device)
            depth_images = depth_images.to(device)
            labels = labels.to(device)

            outputs = model(rgb_images, depth_images)
            _, predicted = torch.max(outputs, 1)

            all_labels.extend(labels.cpu().numpy())
            all_predictions.extend(predicted.cpu().numpy())

    # 计算混淆矩阵
    cm = confusion_matrix(all_labels, all_predictions, labels=np.arange(num_classes))
    return cm

# 调用函数，计算验证集的混淆矩阵
num_classes = 4  # 你的分类数量
cm = calculate_confusion_matrix(val_dataloader, model, device, num_classes)

# 绘制混淆矩阵
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=np.arange(num_classes))
disp.plot(cmap=plt.cm.Blues)
plt.title("Confusion Matrix")
plt.show()
